package cn.bellychang.leetcode.question215;

/**
 * Find the kth largest element in an unsorted array.
 * Note that it is the kth largest element in the sorted order, not the kth distinct element.
 *
 * @author ChangLiang
 * @date 2021/2/2
 */
public class Solution {

    public int findKthLargest(int[] nums, int k) {
        int index = nums.length - k;
        findKthLargest(0, nums.length - 1, nums, index);
        return nums[index];
    }

    private void findKthLargest(int beginIndex, int endIndex, int[] nums, int index) {

        int partitionIndex = partition(beginIndex, endIndex, nums);
        if (partitionIndex != index) {
            if (partitionIndex < index) {
                findKthLargest(partitionIndex + 1, endIndex, nums, index);
            } else {
                findKthLargest(beginIndex, partitionIndex - 1, nums, index);
            }
        }
    }

    private int partition(int beginIndex, int endIndex, int[] nums) {
        int pivot = nums[beginIndex];
        int i = beginIndex + 1;
        int j = beginIndex + 1;
        // loop invariant: [beginIndex, i-1] < pivot  [i, endIndex]剩余部分
        while (j <= endIndex) {
            if (nums[j] < pivot) {
                if (i != j) {
                    swap(i, j, nums);
                }
                i++;
            }
            j++;
        }
        swap(beginIndex, i - 1, nums);
        return i - 1;
    }

    private void swap(int i, int j, int[] nums) {
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }
}
